import keras
import cv2
import numpy as np

# 加载反卷积的训练模型
unet = keras.models.load_model("unet.h5")
# 加载车牌内容识别的模型  图像识别的卷积神经网络
cnn = keras.models.load_model("cnn.h5")

def recognize_panel(image):
    image = image.reshape(1,80,240,3)
    characters = ["京", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑", "苏", "浙", "皖", "闽", "赣", "鲁","豫",
                  "鄂", "湘", "粤", "桂", "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新", "0", "1", "2",
                  "3", "4", "5", "6", "7", "8", "9", "A", "B", "C", "D", "E", "F", "G", "H", "J", "K", "L", "M",
                  "N", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z"]
    result = cnn.predict(image)
    panel_str = ""
    for item in result:
        index = np.argmax(item)
        panel_str = panel_str + characters[index]

    return panel_str[:2] + "." + panel_str[2:]

# -------------------------加载车牌图像，检测车牌区域-----------------------------
img_src_path = "test.jpg"
image = cv2.imdecode(np.fromfile(img_src_path,dtype=np.uint8),-1)
print(image.shape)
if image.shape != (512,512,3):
    image = cv2.resize(image, dsize=(512, 512), interpolation=cv2.INTER_AREA)[:, :, :3]
image = image.reshape(1,512,512,3)

img_mask = unet.predict(image)
img_mask = img_mask.reshape(512,512,3)
img_mask = img_mask / np.max(img_mask) * 255  # 将像素值控制到0- 255 之间
img_mask[:, :, 2] = img_mask[:, :, 1] = img_mask[:, :, 0]  # 三个通道保持相同
img_mask = np.array(img_mask,dtype=np.uint8)

try:
    # opencv3.0
    contours, hierarchy = cv2.findContours(img_mask[:, :, 0], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
except:
    # opencv2.0
    ret, contours, hierarchy = cv2.findContours(img_mask[:, :, 0], cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

image = image.reshape(512,512,3)
for cont in contours:
    x, y, w, h = cv2.boundingRect(cont)
    x0,y0 = x,y
    x1,y1 = x,y + h
    x2,y2 = x + w,y
    x3,y3 = x + w,y + h

    d0, d1, d2, d3 = np.inf,np.inf,np.inf,np.inf
    l0, l1, l2, l3 = (x0, y0), (x1, y1), (x2, y2), (x3, y3)
    for item in cont:
        (current_x,current_y) = item[0]
        dis0 = (current_x - x0)**2 + (current_y - y0)**2
        dis1 = (current_x - x1)**2 + (current_y - y1)**2
        dis2 = (current_x - x2)**2 + (current_y - y2)**2
        dis3 = (current_x - x3)**2 + (current_y - y3)**2
        if dis0 < d0:
            d0 = dis0
            l0 = (current_x,current_y)
        if dis1 < d1:
            d1 = dis1
            l1 = (current_x,current_y)
        if dis2 < d2:
            d2 = dis2
            l2 = (current_x,current_y)
        if dis3 < d3:
            d3 = dis3
            l3 = (current_x,current_y)
    tmp_image = image.copy()
    cv2.line(tmp_image,l0,l2,color=(0,255,0),thickness=2)
    cv2.line(tmp_image,l2,l3,color=(0,255,0),thickness=2)
    cv2.line(tmp_image,l3,l1,color=(0,255,0),thickness=2)
    cv2.line(tmp_image,l1,l0,color=(0,255,0),thickness=2)

    pts1 = np.float32([l0,l2,l1,l3])
    # 车牌的标准大小为240x80
    pst2 = np.float32([
        [0, 0],
        [240, 0],
        [0, 80],
        [240, 80]
    ])

    # 仿射变化的核心
    M = cv2.getPerspectiveTransform(pts1, pst2)
    # 将梯形拉伸成了矩形
    dic_img = cv2.warpPerspective(image, M, (240,80))
    # 识别车牌信息
    result = recognize_panel(dic_img)
    # 绘制到





